const fs = require('fs');
const path = require('path');
const axios = require('axios');
const extract = require('extract-zip');

class DatasetUtils {
    constructor(app) {
        this.app = app;
    }

    async downloadAndInstall(datasetInfo) {
        let dataset_id = datasetInfo.id;
        this.app.queue.clear(`dataset_cloud:${dataset_id}`);

        // 判断下是否存在对应的数据集
        const rows = await this.app.db.datasets.getDatasetInfo(dataset_id);
        if (rows) {
            return { code: -1, msg: `数据集 ${dataset_id} 已经下载！` };
        }

        // 下载数据集并解压
        const dataset_dir = path.join(this.app.userDataPath, 'downloads', dataset_id);
        let result = await this.downloadAndUnzip(dataset_id, datasetInfo.download_url, dataset_dir);
        if (result == false) {
            return { code: -1, msg: `下载解压数据集失败！` };
        }

        // 创建保存list.db的文件夹
        let listdb_dir = path.join(this.app.userDataPath, 'datasets', `${dataset_id}`);
        if (fs.existsSync(listdb_dir)) {
            fs.rmSync(listdb_dir, { recursive: true, force: true });
        }
        fs.mkdirSync(listdb_dir, { recursive: true });

        // 移动list.db文件到对应的目录
        let listdb_path_src = path.join(dataset_dir, 'list.db');
        let listdb_path = path.join(`${listdb_dir}`, 'list.db');
        if (fs.existsSync(listdb_path_src)) {
            fs.renameSync(listdb_path_src, listdb_path);
        } else {
            return { code: -1, msg: '数据集不完整，list.db 未存在！' };
        }

        // 插入数据库记录
        try {
            const lastId = await this.app.db.datasets.insertDataset({
                id: datasetInfo.id,
                datasetName: datasetInfo.dataset_name,
                datasetSuperType: datasetInfo.dataset_supertype,
                datasetType: datasetInfo.dataset_type,
                datasetDir: dataset_dir,
                labelFormat: datasetInfo.label_format,
                totalSamples: datasetInfo.total_samples,
                labeledSamples: datasetInfo.labeled_samples
            });

            let categories_path = path.join(dataset_dir, 'categories.json');
            const categoriesData = fs.readFileSync(categories_path, 'utf-8');
            const categoriesJson = JSON.parse(categoriesData);
            for( let i = 0; i < categoriesJson.length; i++ ){
                let category = categoriesJson[i];
                let categoryInfo = {
                    category_id: category.category_id,
                    category_name: category.category_name,
                    supercategory: category.supercategory
                };
                await this.app.db.categories.insertCategory(dataset_id, categoryInfo, '');
            }

            return { code: 0, msg: '下载并安装数据集成功！' };
        } catch (error) {
            return { code: -1, msg: '插入本地数据集记录失败！' };
        }
    }

    async downloadAndUnzip(dataset_id, url, outputDir) {
        try {
            // Create output directory if it doesn't exist
            if (!fs.existsSync(outputDir)) {
                fs.mkdirSync(outputDir, { recursive: true });
            }
            const tempZipPath = path.join(outputDir, 'download.zip');

            // Start the download
            const response = await axios({
                method: 'GET',
                url: url,
                responseType: 'stream'
            });

            // Get total length from headers
            const totalLength = response.headers['content-length'];

            let total_size = (totalLength / 1024 / 1024).toFixed(2);
            console.log(`Total size: ${total_size} MB`);
            this.app &&
                this.app.queue.enqueue(
                    `dataset_cloud:${dataset_id}`,
                    {
                        event: 'download',
                        data: {
                            total_size
                        }
                    },
                    false
                );

            let downloadedLength = 0;
            let download_percent = 0;
            response.data.on('data', (chunk) => {
                downloadedLength += chunk.length;
                download_percent = ((downloadedLength / totalLength) * 100).toFixed(2);
                process.stdout.write(`Download Progress: ${download_percent}%\r`);

                this.app &&
                    this.app.queue.enqueue(
                        `dataset_cloud:${dataset_id}`,
                        {
                            event: 'download_percent',
                            data: download_percent
                        },
                        false
                    );
            });

            // Stream the file
            const writer = fs.createWriteStream(tempZipPath);
            response.data.pipe(writer);

            await new Promise((resolve, reject) => {
                writer.on('finish', resolve);
                writer.on('error', reject);
            });

            // Use extract-zip with progress
            console.log('\nStarting extraction...');
            let extractedFiles = 0;

            await extract(tempZipPath, {
                dir: outputDir,
                onEntry: (entry) => {
                    extractedFiles++;
                    process.stdout.write(`Extraction Progress: Processed ${extractedFiles} files\r`);
                }
            });

            // Remove the temporary zip file
            fs.unlinkSync(tempZipPath);
            console.log('\nDownload and extraction completed.');
            return true;
        } catch (error) {
            console.error('An error occurred:', error);
            return false;
        }
    }
}

module.exports = { DatasetUtils };
